"""
分层滑点计算器
基于市值和波动率的动态滑点模型
"""

import numpy as np
from typing import Dict, Any
import logging

logger = logging.getLogger(__name__)


class SlippageCalculator:
    """
    分层滑点计算器
    
    特性：
    1. 市值分层：大盘股、中盘股、小盘股
    2. 波动率加权：基于ATR/Price动态调整
    3. 双重保护：默认滑点 + 高波动滑点
    """
    
    def __init__(self, config: Dict[str, Any]):
        """
        初始化滑点计算器
        
        Args:
            config: 回测配置，包含slippage配置
        """
        self.slippage_config = config.get('slippage', {})
        self.market_cap_threshold = None  # 将在需要时传入
        self.high_vol_threshold = self.slippage_config.get('high_volatility_threshold', 0.03)
        
        logger.info(f"滑点计算器初始化完成 - 高波动阈值: {self.high_vol_threshold*100:.1f}%")
    
    def set_market_cap_threshold(self, threshold: Dict[str, float]):
        """设置市值分类阈值"""
        self.market_cap_threshold = threshold
    
    def calculate_atr(self, high: np.ndarray, low: np.ndarray, close: np.ndarray, 
                     period: int = 20) -> float:
        """
        计算平均真实波幅（ATR）
        
        Args:
            high: 最高价序列
            low: 最低价序列
            close: 收盘价序列
            period: 计算周期
        
        Returns:
            ATR值
        """
        if len(high) < period + 1:
            return 0.0
        
        # 计算真实波幅（TR）
        tr1 = high - low  # 当日最高 - 最低
        tr2 = np.abs(high - np.roll(close, 1))  # 当日最高 - 昨收
        tr3 = np.abs(low - np.roll(close, 1))   # 当日最低 - 昨收
        
        # 取最大值作为真实波幅
        tr = np.maximum(tr1, np.maximum(tr2, tr3))
        
        # 计算ATR（移动平均）
        atr = np.mean(tr[-period:])
        
        return atr
    
    def get_stock_type(self, market_cap: float) -> str:
        """
        判断股票类型（大中小盘）
        
        Args:
            market_cap: 流通市值（亿元）
        
        Returns:
            'large', 'mid', 或 'small'
        """
        if self.market_cap_threshold is None:
            logger.warning("市值阈值未设置，使用默认中盘股标准")
            return 'mid'
        
        large_threshold = self.market_cap_threshold.get('large', 200)
        small_threshold = self.market_cap_threshold.get('small', 50)
        
        if market_cap > large_threshold:
            return 'large'
        elif market_cap >= small_threshold:
            return 'mid'
        else:
            return 'small'
    
    def calculate_slippage_rate(self, stock_code: str, current_price: float,
                               hist_data: Dict[str, np.ndarray],
                               market_cap: float = None) -> float:
        """
        计算动态滑点率
        
        Args:
            stock_code: 股票代码
            current_price: 当前价格
            hist_data: 历史数据，包含 'high', 'low', 'close'
            market_cap: 流通市值（亿元），如果为None则使用默认中盘股标准
        
        Returns:
            滑点率（小数形式，如0.001表示0.1%）
        """
        # 1. 确定股票类型
        if market_cap is None or market_cap <= 0:
            stock_type = 'mid'  # 默认中盘股
            logger.debug(f"{stock_code} 市值未知，使用中盘股标准")
        else:
            stock_type = self.get_stock_type(market_cap)
        
        # 2. 获取该类型的滑点配置
        type_config = self.slippage_config.get(f'{stock_type}_cap', {})
        default_rate = type_config.get('default', 0.001)
        high_vol_rate = type_config.get('high_volatility', 0.002)
        
        # 3. 计算波动率
        try:
            high = hist_data.get('high')
            low = hist_data.get('low')
            close = hist_data.get('close')
            
            if high is None or low is None or close is None:
                logger.debug(f"{stock_code} 缺少历史数据，使用默认滑点")
                return default_rate
            
            if len(high) < 21:
                logger.debug(f"{stock_code} 历史数据不足，使用默认滑点")
                return default_rate
            
            # 计算ATR
            atr = self.calculate_atr(high, low, close, period=20)
            
            if atr == 0 or current_price == 0:
                return default_rate
            
            # 计算波动率比率
            volatility_ratio = atr / current_price
            
            # 4. 判断是否为高波动
            if volatility_ratio > self.high_vol_threshold:
                slippage_rate = high_vol_rate
                logger.debug(
                    f"{stock_code} 高波动 (ATR/Price={volatility_ratio*100:.2f}% > {self.high_vol_threshold*100:.1f}%)"
                    f" - 使用{slippage_rate*100:.2f}%滑点"
                )
            else:
                slippage_rate = default_rate
                logger.debug(
                    f"{stock_code} 正常波动 (ATR/Price={volatility_ratio*100:.2f}%)"
                    f" - 使用{slippage_rate*100:.2f}%滑点"
                )
            
            return slippage_rate
            
        except Exception as e:
            logger.warning(f"{stock_code} 计算滑点率异常: {e}，使用默认滑点")
            return default_rate
    
    def apply_slippage(self, price: float, slippage_rate: float, 
                       side: str = 'buy') -> float:
        """
        应用滑点到价格
        
        Args:
            price: 原始价格
            slippage_rate: 滑点率
            side: 'buy' 或 'sell'
        
        Returns:
            应用滑点后的价格
        """
        if side == 'buy':
            # 买入：价格上浮
            return price * (1 + slippage_rate)
        else:
            # 卖出：价格下降
            return price * (1 - slippage_rate)
    
    def get_slippage_summary(self) -> str:
        """
        获取滑点配置摘要
        
        Returns:
            配置摘要字符串
        """
        lines = [
            "=" * 60,
            "分层滑点配置",
            "=" * 60,
            f"高波动阈值: ATR/Price > {self.high_vol_threshold*100:.1f}%",
            "",
            "市值分层滑点:",
        ]
        
        for cap_type in ['large_cap', 'mid_cap', 'small_cap']:
            config = self.slippage_config.get(cap_type, {})
            default = config.get('default', 0) * 100
            high_vol = config.get('high_volatility', 0) * 100
            
            type_name = {
                'large_cap': '大盘股',
                'mid_cap': '中盘股',
                'small_cap': '小盘股'
            }[cap_type]
            
            lines.append(f"  {type_name}: 默认 {default:.2f}%, 高波动 {high_vol:.2f}%")
        
        lines.append("=" * 60)
        return "\n".join(lines)


if __name__ == '__main__':
    # 测试滑点计算器
    import logging
    logging.basicConfig(level=logging.DEBUG)
    
    # 模拟配置
    test_config = {
        'slippage': {
            'high_volatility_threshold': 0.03,
            'large_cap': {'default': 0.0005, 'high_volatility': 0.001},
            'mid_cap': {'default': 0.001, 'high_volatility': 0.002},
            'small_cap': {'default': 0.002, 'high_volatility': 0.005}
        }
    }
    
    calculator = SlippageCalculator(test_config)
    calculator.set_market_cap_threshold({'large': 200, 'small': 50})
    
    print(calculator.get_slippage_summary())
    
    # 测试1：低波动中盘股
    print("\n" + "=" * 60)
    print("测试1: 低波动中盘股")
    hist_data = {
        'high': np.array([10.5, 10.6, 10.4, 10.5, 10.6] * 5),
        'low': np.array([10.3, 10.4, 10.2, 10.3, 10.4] * 5),
        'close': np.array([10.4, 10.5, 10.3, 10.4, 10.5] * 5)
    }
    rate = calculator.calculate_slippage_rate('000001.SZ', 10.5, hist_data, market_cap=100)
    print(f"滑点率: {rate*100:.2f}%")
    print(f"买入价: {calculator.apply_slippage(10.5, rate, 'buy'):.4f}")
    print(f"卖出价: {calculator.apply_slippage(10.5, rate, 'sell'):.4f}")
    
    # 测试2：高波动小盘股
    print("\n" + "=" * 60)
    print("测试2: 高波动小盘股")
    hist_data_volatile = {
        'high': np.array([10.8, 11.2, 10.6, 11.0, 10.9] * 5),
        'low': np.array([10.0, 10.2, 9.8, 10.1, 10.0] * 5),
        'close': np.array([10.4, 10.8, 10.2, 10.5, 10.4] * 5)
    }
    rate = calculator.calculate_slippage_rate('300001.SZ', 10.5, hist_data_volatile, market_cap=30)
    print(f"滑点率: {rate*100:.2f}%")
    print(f"买入价: {calculator.apply_slippage(10.5, rate, 'buy'):.4f}")
    print(f"卖出价: {calculator.apply_slippage(10.5, rate, 'sell'):.4f}")

